from shell.knowledge_graph_mining.prepare_graph_songs import GraphSongs
import json
from node2vec import Node2Vec

# Create a graph 这里可以给出自己的graph
json_str = json.load(open('../../data/data.json'))
gs_obj = GraphSongs(json_str)
graph = gs_obj.prepare_graph(list(json_str.keys()))

node2vec = Node2Vec(graph, dimensions=64, walk_length=10, num_walks=10, workers=4)

model = node2vec.fit(workers=4)
# most similar nodes
# res = model.wv.most_similar('最浪漫的事')
# print(res)

model.save('graph_embedding.bin')
